#pragma once

#include <thrust/transform.h>

template <typename Container1,
          typename Container2,
          typename UnaryOperation>
struct TransformUnaryOp
{
  Container1 A;
  Container2 B;
  UnaryOperation unary_op;

  template <typename Range1, typename Range2>
  TransformUnaryOp(const Range1& X, const Range2& Y, UnaryOperation unary_op)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      unary_op(unary_op)
  {}

  void operator()(void)
  {
    thrust::transform(A.begin(), A.end(), B.begin(), unary_op);
  }
  
  void reset(void)
  {
    // nothing to do
  }
};

template <typename Container1,
          typename Container2,
          typename Container3,
          typename UnaryOperation,
          typename Predicate>
struct TransformIfUnaryOp
{
  Container1 A; // input
  Container2 B; // stencil
  Container3 C; // output
  UnaryOperation unary_op;
  Predicate pred;

  template <typename Range1, typename Range2, typename Range3>
  TransformIfUnaryOp(const Range1& X, const Range2& Y, const Range3& Z, UnaryOperation unary_op, Predicate pred)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      C(Z.begin(), Z.end()),
      unary_op(unary_op),
      pred(pred)
  {}

  void operator()(void)
  {
    thrust::transform_if(A.begin(), A.end(), B.begin(), C.begin(), unary_op, pred);
  }
  
  void reset(void)
  {
    // nothing to do
  }
};


template <typename Container1,
          typename Container2,
          typename Container3,
          typename BinaryOperation>
struct TransformBinaryOp
{
  Container1 A;
  Container2 B;
  Container3 C;
  BinaryOperation binary_op;

  template <typename Range1, typename Range2, typename Range3>
  TransformBinaryOp(const Range1& X, const Range2& Y, const Range3& Z, BinaryOperation binary_op)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      C(Z.begin(), Z.end()),
      binary_op(binary_op)
  {}

  void operator()(void)
  {
    thrust::transform(A.begin(), A.end(), B.begin(), C.begin(), binary_op);
  }
  
  void reset(void)
  {
    // nothing to do
  }
};


template <typename Container1,
          typename Container2,
          typename Container3,
          typename Container4,
          typename BinaryOperation,
          typename Predicate>
struct TransformIfBinaryOp
{
  Container1 A; // input
  Container2 B; // input
  Container3 C; // stencil
  Container4 D; // output
  BinaryOperation binary_op;
  Predicate pred;

  template <typename Range1, typename Range2, typename Range3, typename Range4>
  TransformIfBinaryOp(const Range1& X, const Range2& Y, const Range3& Z, const Range4& W, BinaryOperation binary_op, Predicate pred)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      C(Z.begin(), Z.end()),
      D(W.begin(), W.end()),
      binary_op(binary_op),
      pred(pred)
  {}

  void operator()(void)
  {
    thrust::transform_if(A.begin(), A.end(), B.begin(), C.begin(), D.begin(), binary_op, pred);
  }
  
  void reset(void)
  {
    // nothing to do
  }
};

